import config


def getNet_BraTs(net_name, g_config=config):
    if net_name == 'SwinUnet':
        import nets.SwinUnet.networks.swin_transformer_unet_skip_expand_decoder_sys as net
        model = net.__dict__['SwinTransformerSys'](img_size=160,
                                                   patch_size=4,
                                                   window_size=5)

    elif net_name == 'SwinUnet_ps2':
        import nets.SwinUnet.networks.swin_transformer_unet_skip_expand_decoder_sys as net
        model = net.__dict__['SwinTransformerSys'](img_size=160,
                                                   patch_size=2,
                                                   window_size=5)

    elif net_name == 'SwinUnet_ps2_plus1layer':
        import nets.SwinUnet.networks.swin_transformer_unet_skip_expand_decoder_sys as net
        model = net.__dict__['SwinTransformerSys'](img_size=160,
                                                   patch_size=2,
                                                   window_size=5,
                                                   depths=[2, 2, 2, 2, 2],
                                                   depths_decoder=[1, 2, 2, 2, 2],
                                                   num_heads=[3, 6, 12, 24, 48])

    elif net_name == 'SwinUnet_pt':
        import nets.SwinUnet.networks.swin_unet_custom as net
        model = net.__dict__['SwinUnet']()

    elif net_name == 'UNet':
        import nets.UNet.UNet as net
        model = net.__dict__['UNet']()

    elif net_name == 'UNet1':
        import nets.UNet1.UNet as net
        model = net.__dict__['UNet']()

    elif net_name == 'TransUNet':
        import nets.TransUNet.networks.vit_seg_modeling as net
        model = net.__dict__['VisionTransformer']()

    elif net_name == 'TransFuse':
        import nets.TransFuse.lib.TransFuse as net
        model = net.__dict__['TransFuse_S'](num_classes=g_config.num_classes)

    elif net_name == 'TransBTS':
        import nets.TransBTS.models.TransBTS.TransBTS_downsample8x_skipconnection as net
        model = net.__dict__['BTS'](img_dim=128,
                                    patch_dim=8,
                                    num_channels=g_config.input_classes,
                                    num_classes=g_config.num_classes,
                                    embedding_dim=512,
                                    num_heads=8,
                                    num_layers=4,
                                    hidden_dim=4096,
                                    dropout_rate=0.1,
                                    attn_dropout_rate=0.1,
                                    conv_patch_representation=True,
                                    positional_encoding_type="learned",
                                    )

    elif net_name == 'SwinUnetwDAmoduleV1':
        import nets.SwinUnetwDAmodule.SwinUnetwDAmoduleV3 as net
        model = net.__dict__['SwinTransformerSys'](img_size=160,
                                                   patch_size=2,
                                                   window_size=5,
                                                   depths=[2, 2, 2, 2, 2],
                                                   depths_decoder=[1, 2, 2, 2, 2],
                                                   num_heads=[3, 6, 12, 24, 48])

    elif net_name == 'SwinUnetwDAmoduleV2':
        import nets.SwinUnetwDAmodule.SwinUnetwDAmoduleV2 as net
        model = net.__dict__['SwinTransformerSys'](img_size=160,
                                                   patch_size=2,
                                                   window_size=5,
                                                   depths=[2, 2, 2, 2, 2],
                                                   depths_decoder=[1, 2, 2, 2, 2],
                                                   num_heads=[3, 6, 12, 24, 48])

    elif net_name == 'SwinUnetwDAmoduleV3':
        import nets.SwinUnetwDAmodule.SwinUnetwDAmoduleV3 as net
        model = net.__dict__['SwinTransformerSys'](img_size=160,
                                                   patch_size=2,
                                                   window_size=5,
                                                   depths=[2, 2, 2, 2, 2],
                                                   depths_decoder=[1, 2, 2, 2, 2],
                                                   num_heads=[3, 6, 12, 24, 48])

    elif net_name == 'SwinUnetwNAmodule':
        import nets.SwinUnetwNAmodule.SwinUnetwNAm as net
        model = net.__dict__['SwinTransformerSys'](img_size=160,
                                                   patch_size=2,
                                                   window_size=5,
                                                   depths=[2, 2, 2, 2, 2],
                                                   depths_decoder=[1, 2, 2, 2, 2],
                                                   num_heads=[3, 6, 12, 24, 48])

    elif net_name == 'SwinUnetwLightViTBlock':
        import nets.SwinUnetwLightViTBlock.SwinUnetwLightViTBlock as net
        model = net.__dict__['SwinTransformerSys'](img_size=160,
                                                   patch_size=2,
                                                   window_size=5,
                                                   depths=[2, 2, 2],
                                                   depths_decoder=[1, 2, 2],
                                                   num_heads=[3, 6, 12])

    elif net_name == 'LightViT':
        import nets.LightViT.lightvit as net
        model = net.__dict__['LightViT'](img_size=160,
                                         in_chans=g_config.input_classes,
                                         num_classes=g_config.num_classes,
                                         patch_size=2,
                                         window_size=5,
                                         num_layers=[2, 2, 2, 2, 2],
                                         embed_dims=[32, 64, 160, 256, 256],
                                         num_heads=[1, 2, 5, 8, 8],
                                         mlp_ratios=[8, 8, 4, 4, 4])

    else:
        pass
        return None
    return model


def getNet_Synapse(net_name, g_config=config):
    if net_name == 'SwinUnet':
        import nets.SwinUnet.networks.swin_transformer_unet_skip_expand_decoder_sys as net
        model = net.__dict__['SwinTransformerSys'](img_size=224,
                                                   patch_size=4,
                                                   window_size=g_config.window_size,
                                                   in_chans=g_config.input_classes_synapse,
                                                   num_classes=g_config.num_classes_synapse)

    elif net_name == 'SwinUnet_ps2':
        import nets.SwinUnet.networks.swin_transformer_unet_skip_expand_decoder_sys as net
        model = net.__dict__['SwinTransformerSys'](img_size=160,
                                                   patch_size=2,
                                                   window_size=5)

    elif net_name == 'SwinUnet_ps2_plus1layer':
        import nets.SwinUnet.networks.swin_transformer_unet_skip_expand_decoder_sys as net
        model = net.__dict__['SwinTransformerSys'](img_size=224,
                                                   patch_size=2,
                                                   window_size=5,
                                                   depths=[2, 2, 2, 2, 2],
                                                   depths_decoder=[1, 2, 2, 2, 2],
                                                   num_heads=[3, 6, 12, 24, 48])

    elif net_name == 'SwinUnet_pt':
        import nets.SwinUnet.networks.swin_unet_custom as net
        model = net.__dict__['SwinUnet']()

    elif net_name == 'UNet':
        import nets.UNet.UNet as net
        model = net.__dict__['UNet']()

    elif net_name == 'UNet1':
        import nets.UNet1.UNet as net
        model = net.__dict__['UNet'](in_chan=g_config.input_classes_synapse,
                                     out_chan=g_config.num_classes_synapse)

    elif net_name == 'TransUNet':
        import nets.TransUNet.networks.vit_seg_modeling as net
        model = net.__dict__['VisionTransformer']()

    elif net_name == 'TransFuse':
        import nets.TransFuse.lib.TransFuse as net
        model = net.__dict__['TransFuse_S'](num_classes=g_config.num_classes)

    elif net_name == 'TransBTS':
        import nets.TransBTS.models.TransBTS.TransBTS_downsample8x_skipconnection as net
        model = net.__dict__['BTS'](img_dim=128,
                                    patch_dim=8,
                                    num_channels=g_config.input_classes,
                                    num_classes=g_config.num_classes,
                                    embedding_dim=512,
                                    num_heads=8,
                                    num_layers=4,
                                    hidden_dim=4096,
                                    dropout_rate=0.1,
                                    attn_dropout_rate=0.1,
                                    conv_patch_representation=True,
                                    positional_encoding_type="learned",
                                    )

    elif net_name == 'SwinUnetwDAmoduleV1':
        import nets.SwinUnetwDAmodule.SwinUnetwDAmoduleV3 as net
        model = net.__dict__['SwinTransformerSys'](img_size=160,
                                                   patch_size=2,
                                                   window_size=5,
                                                   depths=[2, 2, 2, 2, 2],
                                                   depths_decoder=[1, 2, 2, 2, 2],
                                                   num_heads=[3, 6, 12, 24, 48])

    elif net_name == 'SwinUnetwDAmoduleV2':
        import nets.SwinUnetwDAmodule.SwinUnetwDAmoduleV2 as net
        model = net.__dict__['SwinTransformerSys'](img_size=160,
                                                   patch_size=2,
                                                   window_size=5,
                                                   depths=[2, 2, 2, 2, 2],
                                                   depths_decoder=[1, 2, 2, 2, 2],
                                                   num_heads=[3, 6, 12, 24, 48])

    elif net_name == 'SwinUnetwDAmoduleV3':
        import nets.SwinUnetwDAmodule.SwinUnetwDAmoduleV3 as net
        model = net.__dict__['SwinTransformerSys'](img_size=160,
                                                   patch_size=2,
                                                   window_size=5,
                                                   depths=[2, 2, 2, 2, 2],
                                                   depths_decoder=[1, 2, 2, 2, 2],
                                                   num_heads=[3, 6, 12, 24, 48])

    elif net_name == 'SwinUnetwNAmodule':
        import nets.SwinUnetwNAmodule.SwinUnetwNAm as net
        model = net.__dict__['SwinTransformerSys'](img_size=224,
                                                   patch_size=4,
                                                   window_size=g_config.window_size,
                                                   num_heads=[3, 6, 12, 24],
                                                   in_chans=g_config.input_classes_synapse,
                                                   num_classes=g_config.num_classes_synapse)

    else:
        pass
        return None
    return model
